import asyncio
import websockets
import json
import logging
import numpy as np
from typing import Dict, Callable
from ws_comm import Packet
logger = logging.getLogger("model_infer_ws")
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
class ModelInferWebSocket:
    def __init__(self, 
                 get_obs_callback: Callable[[], Dict],
                 take_action_callback: Callable[[np.ndarray], Dict],
                 host: str = "127.0.0.1",
                 port: int = 8765):
        self.get_obs_callback = get_obs_callback
        self.take_action_callback = take_action_callback
        self.host = host
        self.port = port
        self.url = f"ws://{host}:{port}"
        self.websocket = None
        self.running = False
    async def handle_packet(self, packet: Packet) -> Packet:
        packet_dict = packet.to_dict()
        packet_type = packet_dict["type"]
        if packet_type == "GET_OBS":
            logger.info("收到 GET_OBS 请求")
            try:
                obs_dict = self.get_obs_callback()
                obs_serialized = self.serialize_observation(obs_dict)
                response_data = {
                    "obs": obs_serialized,
                    "success": True,
                    "error": None
                }
            except Exception as e:
                logger.error(f"获取观察失败: {e}", exc_info=True)
                response_data = {
                    "obs": None,
                    "success": False,
                    "error": str(e)
                }
            return Packet("GET_OBS_RESPONSE", response_data)
        elif packet_type == "TAKE_ACTION":
            logger.info(f"收到 TAKE_ACTION 请求")
            try:
                action_data = packet_dict["data"]["action"]
                action_array = np.array(action_data, dtype=np.float32)
                result = self.take_action_callback(action_array)
                response_data = {
                    "success": result.get("success", True),
                    "error": result.get("error", None)
                }
            except Exception as e:
                logger.error(f"执行动作失败: {e}", exc_info=True)
                response_data = {
                    "success": False,
                    "error": str(e)
                }
            return Packet("TAKE_ACTION_RESPONSE", response_data)
        else:
            logger.warning(f"未知的包类型: {packet_type}")
            return Packet("ERROR_RESPONSE", {"success": False, "error": f"Unknown packet type: {packet_type}"})
    def serialize_observation(self, obs_dict: Dict) -> Dict:
        serialized = {}
        for key, value in obs_dict.items():
            if isinstance(value, np.ndarray):
                serialized[key] = value.tolist()
            elif isinstance(value, dict):
                serialized[key] = self.serialize_observation(value)
            elif isinstance(value, (list, tuple)):
                serialized[key] = [self.serialize_observation(item) if isinstance(item, dict) else (item.tolist() if isinstance(item, np.ndarray) else item) for item in value]
            else:
                serialized[key] = value
        return serialized
    async def client_handler(self, websocket):
        self.websocket = websocket
        logger.info(f"已连接到云端服务器: {websocket.remote_address}")
        try:
            async for message in websocket:
                try:
                    packet_dict = json.loads(message)
                    packet = Packet.from_dict(packet_dict)
                    response_packet = await self.handle_packet(packet)
                    response_str = json.dumps(response_packet.to_dict(), default=str)
                    await websocket.send(response_str)
                except json.JSONDecodeError as e:
                    logger.error(f"JSON 解析失败: {e}")
                    error_response = Packet("ERROR_RESPONSE", {"success": False, "error": f"JSON decode error: {e}"})
                    await websocket.send(json.dumps(error_response.to_dict(), default=str))
                except Exception as e:
                    logger.error(f"处理消息失败: {e}", exc_info=True)
                    error_response = Packet("ERROR_RESPONSE", {"success": False, "error": str(e)})
                    await websocket.send(json.dumps(error_response.to_dict(), default=str))
        except websockets.exceptions.ConnectionClosed:
            logger.info("与云端服务器的连接已关闭")
        except Exception as e:
            logger.error(f"WebSocket 连接错误: {e}", exc_info=True)
        finally:
            self.websocket = None
    async def connect_and_run(self):
        self.running = True
        while self.running:
            try:
                logger.info(f"正在连接到云端服务器: {self.url}")
                async with websockets.connect(self.url, max_size=None) as websocket:
                    await self.client_handler(websocket)
            except (ConnectionRefusedError, OSError) as e:
                logger.warning(f"连接被拒绝或网络错误: {e}，5秒后重试...")
                await asyncio.sleep(5)
            except websockets.exceptions.InvalidMessage as e:
                logger.error(f"无效的 HTTP 响应，可能服务器未正确运行或端口转发有问题: {e}")
                logger.info("请检查：1. SSH 端口转发是否正常 2. 云端服务器是否运行 3. 端口是否正确")
                await asyncio.sleep(5)
            except websockets.exceptions.InvalidStatusCode as e:
                logger.error(f"WebSocket 握手失败，状态码无效: {e}，5秒后重试...")
                await asyncio.sleep(5)
            except Exception as e:
                logger.error(f"连接错误: {type(e).__name__}: {e}，5秒后重试...")
                await asyncio.sleep(5)
    def stop(self):
        self.running = False
        if self.websocket:
            asyncio.create_task(self.websocket.close())
    async def run(self):
        await self.connect_and_run()
def default_get_obs() -> Dict:
    dummy_image = np.zeros((240, 320, 3), dtype=np.uint8)
    dummy_pointcloud = np.zeros((1000, 3), dtype=np.float32)
    dummy_joint_state = np.zeros(14, dtype=np.float32)
    dummy_ee_state = np.zeros(16, dtype=np.float32)
    return {
        "image": dummy_image,
        "point_cloud": dummy_pointcloud,
        "joint_state": dummy_joint_state,
        "ee_state": dummy_ee_state
    }
def default_take_action(action: np.ndarray) -> Dict:
    logger.info(f"执行动作: shape={action.shape}, action={action.tolist()}")
    return {"success": True, "error": None}
if __name__ == "__main__":
    model_infer_ws = ModelInferWebSocket(
        get_obs_callback=default_get_obs,
        take_action_callback=default_take_action,
        host="127.0.0.1",
        port=8765
    )
    try:
        asyncio.run(model_infer_ws.run())
    except KeyboardInterrupt:
        logger.info("接收到中断信号，正在关闭...")
        model_infer_ws.stop()

